import json
import os

with open('bladder.json', 'r', encoding='utf-8') as f:
    data = json.load(f)

with open('../bladder/bladder1204_train.json', 'r', encoding='utf-8') as f:
    e = json.load(f)
# with open('../bladder/bladder_val.json', 'r', encoding='utf-8') as f:
#     val = json.load(f)
#
# with open('../bladder/bladder_test.json', 'r', encoding='utf-8') as f:
#     test = json.load(f)

e_list = [info['imageId'] for info in e]
# test_name = [info['imageId'] for info in test]
# print(e_name)

bladder_train = 'bladder1204_train.json'
# bladder_val = 'bladder_val.json'
# bladder_test = 'bladder_test.json'
bladder_train_list, bladder_val_list, bladder_test_list = [], [], []
for info in data:
    if info['imageId'] in e_list:
        if info['type'] == 'noprolapsed' or info['type'] == 'mild':
            bladder_val_list.append(info)
        else:
            info['type'] = 'serious'
            bladder_val_list.append(info)
    # if info['imageId'] in test_name:
    #     if info['type'] == 'noprolapsed' or info['type'] == 'mild':
    #         bladder_test_list.append(info)
    #     else:
    #         info['type'] = 'serious'
    #         bladder_test_list.append(info)

with open(bladder_train, 'w', encoding='utf-8') as f:
    json.dump(bladder_val_list, f)

# with open(bladder_val, 'w', encoding='utf-8') as f:
#     json.dump(bladder_val_list, f)
#
# with open(bladder_test, 'w', encoding='utf-8') as f:
#     json.dump(bladder_test_list, f)
